# 该程序是将DOT设备的轮廓图对比度增强的程序，拿到对比度增强的结果就可以进行推理
# 搭建推理环境7
# conda create -n yolov8 python=3.8
# conda install pytorch==1.11.0 torchvision==0.12.0 torchaudio==0.11.0 cpuonly -c pytorch

import cv2
import numpy as np
import os 
from ultralytics import YOLO
import scipy.io
import numpy as np
import matplotlib.pyplot as plt
import csv
from scipy.interpolate import interp1d
from matplotlib.path import Path
import matplotlib.image as mpimg


class boundary:
    # fence_fine_tuning csv的围挡比轮廓中宽，所以得调整宽高一样再拉伸，fence_fine_tuning可以认为二者轮廓差，单位mm
    # translation 平移轮廓单位mm
    def __init__(self, csvfile ,PF_image,fence_fine_tuning=[20,10],translation=[15]):
        # config 不通过参数暴露
        self.temp_dir = 'temp' #中间结果保存目录
        self.sample_count = 30 #采样点个数用于平滑
        self.confidence = 0.75
        self.mask_out_dir = 'mask_out_dir'
        #############################
        self.csv_file_name = csvfile
        self.PF_image_name = PF_image
        self.fence_fine_tuning = fence_fine_tuning
        self.translation = translation
        self.opencv_boundary = []
        self.yolo_boundary = []
        self.height = 0 #轮廓被拉伸到640不加padding的宽高
        self.width = 0  
        self.csv_height = 0 #轮廓被拉伸到640不加padding的宽高
        self.csv_width = 0  
        self.padding = 0
        self.points = None # 原始轮廓点
        
        
    
    # 读取原始图片加入拉伸640*640供推理端使用，如果不用推理则不需要
    def read_image_with_padding(self,image_file):
        image = cv2.imread(image_file, cv2.IMREAD_GRAYSCALE)
        # 获取图像的高度和宽度
        height, width = image.shape
        image_scale  = 0
        # 计算调整后的尺寸
        if width > height:
            new_width = 640
            new_height = int(height * (new_width / width))
            image_scale = new_width / width
        else:
            new_height = 640
            new_width = int(width * (new_height / height))
            image_scale = new_height / height
        self.height = new_height
        self.width = new_width
        # 按原比例拉伸到640*640
        resized_image_640 = np.zeros((640, 640), dtype=np.uint8)
        resized_image_640_enhance = np.zeros((640, 640), dtype=np.uint8)
        padding = int((640 - new_width)/2)
        self.padding = padding
        resized_image_640[:new_height, padding:new_width+padding] = cv2.resize(image, (new_width, new_height))
        return resized_image_640
    
    # 推理轮廓并减去padding
    def  inference_removed_padding(self,image_file=None):
            if image_file is None:
                image_file = self.PF_image_name
            x_values_remove_padding =[]
            y_values =[]
            resized_image_640 = self.read_image_with_padding(image_file)
            resized_image_640_enhance = self.enhance_strategy1(resized_image_640)
            file_name = os.path.basename(image_file)
            # file_name = os.path.splitext(os.path.basename(image_file))[0]
            # 输出文件保存临时增强后的图像，用于推理
            if not os.path.exists(self.temp_dir):
            # 创建文件夹
                os.makedirs(self.temp_dir)
            output_path = os.path.join(self.temp_dir, file_name)
            cv2.imwrite(output_path, resized_image_640_enhance)
            model = YOLO(r'.\best.pt') #加载模型
            results=model.predict(output_path, conf=self.confidence ,save=True,mask_ratio=1,retina_masks=True,max_det=1,save_txt=True,device='cpu') #进行推理
            if results[0].masks is not None:
                masks0 = results[0].masks.xy
                masks0 = masks0[0]    
                coordinates = [[int(x), int(y)] for x, y in masks0]
                x_values = [int(x[0]) for x in coordinates]
                y_values = [int(x[1]) for x in coordinates]
        
            # # # debug
            # points = np.column_stack((x_values, y_values))
            # cv2.drawContours(resized_image_640_enhance, [points], 0, (0, 255, 255), 2)
            # cv2.imshow('Image with Labels', resized_image_640_enhance)
            # cv2.waitKey(0)
            # cv2.destroyAllWindows()    
        
            # 还原points的坐标，移除padding操作带来的影响
            padding = self.padding
            x_values_remove_padding = [x-padding for x in x_values]
            x_values_remove_padding = [0 if num < 0 else num for num in x_values_remove_padding] #把小于0的数变为0

            # 还原原始图像上的像素位置，像素得取整         
            x_values_remove_padding = [int(x) for x in x_values_remove_padding]
            y_values = [int(y) for y in y_values]
            self.points = np.column_stack((x_values_remove_padding, y_values))
            # # # debug    
            # plt.scatter(x_values_remove_padding, y_values, color='red')
            # plt.show() 
        

    
    # 直方图均衡
    def enhance_strategy1(self,image):
        equalized_row_image = np.zeros((image.shape[0], image.shape[1]), dtype=np.uint8)
        # 对每一行应用直方图均衡
        for i in range(image.shape[0]):
            equalized_row = cv2.equalizeHist(np.expand_dims(image[i, :], axis=0))
            equalized_row_image[i, :] = equalized_row
        # cv2.imshow('Image with Labels', equalized_row_image)
        # cv2.waitKey(0)
        # cv2.destroyAllWindows() 
        # return equalized_row_image 

    
    # 对比度均衡
    def enhance_strategy2(self,image):  
        #TODO  
        pass
    
    # 轮廓旋转
    def points_rotate(self,L_R,points=None):    
        if points is None:
            points = self.points
        # 坐标转置
        swapped_points = points[:, ::-1]
        # 提取 x 和 y 坐标
        x = swapped_points[:, 0]
        y = swapped_points[:, 1]
        if L_R == 'R': # 右侧旋转方式
            max_y = max(y)
            y = -y+max_y        
        if L_R == 'L': # 左侧旋转方式
            max_x = max(x)
            x = -x+max_x
        self.points = np.column_stack((x, y))
        # 更新旋转后的高宽
        self.height, self.width= self.width,self.height
        # plt.scatter(x, y, color='red')
        # plt.gca().invert_yaxis()
        # plt.show() 
            
    # 轮廓拉伸 stretch
    def points_stretch(self,points=None):     
        if points is None:
            points = self.points
        csv_data = [] 
        csv_num_rows = 0
        csv_num_columns = 0
        # 打开 CSV 文件
        with open(csv_file, 'r') as file:
            # 创建 CSV 读取器
            reader = csv.reader(file)
            header = next(reader)
            csv_num_columns = len(header)
            csv_num_rows = sum(1 for _ in reader)
            self.csv_height = csv_num_rows
            self.csv_width = csv_num_columns
            # # 打印行列信息
            # print(f"行数: {num_rows}")
            # print(f"列数: {num_columns}")
        with open(csv_file, 'r') as file:
            reader = csv.reader(file)
            for row in reader:
                csv_data.append([float(val)  for val in row if val != ""])
            csv_data = np.array(csv_data) 
        # 拉伸变换
        scale_h = (csv_num_rows-(self.fence_fine_tuning[1]/0.156))/self.height    
        scale_w = (csv_num_columns-(self.fence_fine_tuning[0]/0.156))/self.width
        #plt.scatter(points[:, 0], points[:, 1], color='red')
        points[:, 0] = scale_w*points[:, 0]
        points[:, 1] = scale_h*points[:, 1]
        points[:, 0] = points[:, 0]+self.translation[0]/0.156 #右移 
        self.points = points
        # plt.scatter(points[:, 0], points[:, 1], color='blue')
        # plt.gca().invert_yaxis()
        # plt.show()
        
    
    # 轮廓平滑过滤
    def boundary_smoothing(self,points=None):
        if points is None:
            points = self.points
        # epsilon = 0.005 * cv2.arcLength(points_results, True)
        # approx = cv2.approxPolyDP(points_results, epsilon, True)
        # return approx
        x = points[:, 0]
        y = points[:, 1]
        filtered_x = []
        filtered_y = []
        for x_val, y_val in zip(x, y):
            if x_val >= 0 and y_val >= 0:
                filtered_x.append(x_val)
                filtered_y.append(y_val)
        #len_x = len(filtered_x)   
        # 将 filtered_x 和 filtered_y 打包并根据 filtered_x 的值进行排序
        sorted_data = sorted(zip(filtered_x, filtered_y))
        # 解压排序后的数据回 filtered_x 和 filtered_y
        filtered_x, filtered_y = zip(*sorted_data)
        sampled_x = list(filtered_x[0::self.sample_count])
        sampled_y = list(filtered_y[0::self.sample_count])
        interval = int((filtered_x[-1]-filtered_x[0])/self.sample_count)
        if sampled_y[0]!=0:
            sampled_x.insert(0, sampled_x[0]-interval)
            sampled_y.insert(0,0)
        if sampled_y[-1]!=0:
            sampled_x.insert(len(sampled_x), sampled_x[-1]+interval)
            sampled_y.insert(len(sampled_y),0)
        interp_func = interp1d(sampled_x, sampled_y, kind='cubic')
        num_points = int(max(sampled_x) - min(sampled_x))   # 插值点的数量
        x_interp = np.linspace(min(sampled_x), max(sampled_x), num_points)
        y_interp = interp_func(x_interp)
        self.points = np.column_stack((x_interp, y_interp))
        # plt.scatter(points[:, 0], points[:, 1], color='red')
        # plt.scatter(x_interp, y_interp, color='blue')
        # plt.gca().invert_yaxis()
        # plt.show()
        
    def opencv_backend(self):
        #TODO 
        pass
    
    def yolo_backend(self,L_R):
        #TODO
        self.read_image_with_padding(self.PF_image_name)
        self.inference_removed_padding() 
        self.points_rotate(L_R)
        self.points_stretch()
        self.boundary_smoothing()
    
    def show_image(self,image,points):
        #TODO
        pass


    def save_mask(self):
        csv_num_columns = self.csv_width
        csv_num_rows = self.csv_height
        image = np.zeros((csv_num_rows, csv_num_columns,4),dtype=np.uint8)
        path = Path(self.points)
        x = np.arange(csv_num_columns)
        y = np.arange(csv_num_rows)
        xv, yv = np.meshgrid(x, y)
        mask = path.contains_points(np.column_stack((xv.flatten(), yv.flatten())))
        mask = mask.reshape((csv_num_rows, csv_num_columns))
        image[mask] = [0, 0, 0, 0]
        image[~mask] = [255, 255, 255, 255]
        # 在当前路径保存掩模mask.png   
        if not os.path.exists(self.mask_out_dir):
            # 创建文件夹
            os.makedirs(self.mask_out_dir)
        absolute_path = os.path.join(self.mask_out_dir, 'mask.png')
        mpimg.imsave(absolute_path, image, cmap='gray')

    def find_breast_boundary(self,csv_file=None,image_file=None):
        if csv_file is None:
            csv_file = self.csv_file_name
        if image_file is None:
            image_file = self.PF_image_name
        imagename = os.path.basename(input_image)
        if 'R' in imagename:
            self.yolo_backend('R')
        elif 'L' in imagename:
            self.yolo_backend('L')
        else :
            error_message = 'There is have not mask.png'
            raise FileNotFoundError(error_message) 

if __name__ == "__main__":
    
    # 这些路径可以自定义根据自己情况
    input_image = r'E:\code\imagedata\上海四院\202312210006\origin\202312210006_L_A_P.png' # 输入的轮廓原始图
    csv_file = r'E:\code\imagedata\上海四院\202312210006\L\L_1.csv' # 对应的csv文件
    mask_outpath = r'.' # 最后掩模生成的路径
    # 附带测试用例
    # 乳房左侧202311280001_L_PF.png
    # 乳房右侧202311210002_R_PF.png
    # 1702968918645.png 推不出轮廓时候抛出异常
    # 接口函数
    a = boundary(csv_file,input_image)
    a.find_breast_boundary()
    a.save_mask()